#!/usr/bin/python

import sys, os, re, random

if len(sys.argv) < 3:
        print "Usage:"
        print "\t python log-watcher.py <iptables_filename> <eps_filename> [<title>]"
        sys.exit(2)

try:
	fin = open(sys.argv[1], 'r')
except:
	print "Cannot open %s file" % sys.argv[1]
	sys.exit(2)

if len(sys.argv) == 4:
	title = sys.argv[3]
else:
	title = "Port activity"

hash_port = {}
hash_dst = {}

for line in fin.readlines():
	pattern = re.search("SRC=([^ ]*).*DST=([^ ]*).*SPT=([^ ]*).*DPT=([^ ]*)", line);

	try:
		src = pattern.group(1)
		dst = pattern.group(2)
		dst_host = int(re.search("(.*)\.(.*)\.(.*)\.(.*)", dst).group(4))
		spt = int(pattern.group(3))
		dpt = int(pattern.group(4))	
	except:
		""" nothing to do """

	# the power of python ;)
	if dpt in hash_port:
		if (src, dst) not in hash_port[dpt]:
			hash_port[dpt][(src, dst)] = True
	else:
		hash_port[dpt] = {}
		hash_port[dpt][(src, dst)] = True
	
	if dst_host in hash_dst:
		if src not in hash_dst[dst_host]:
			hash_dst[dst_host][src] = True
	else:
		hash_dst[dst_host] = {}
		hash_dst[dst_host][src] = True

ports = hash_port.keys()
ports.sort()

dsts = hash_dst.keys()
dsts.sort()

# 2 files to be used by gnuplot
tmp_fname1 = "log-watcher" + str(random.random() * 10000)
tmp_fname2 = "log-watcher" + str(random.random() * 10000)

# the same for destinations graph
tmp_fname1_dst = "log-watcher" + str(random.random() * 10000)
tmp_fname2_dst = "log-watcher" + str(random.random() * 10000)

try:
	gnuplot_in = open(tmp_fname1, "w")
except:
	print "Cannot create tmp file %s" % tmp_fname1
	sys.exit(2)

try:
	gnuplot = open(tmp_fname2, "w")
except:
	print "Cannot create tmp file %s" % tmp_fname2
	sys.exit(2)

try:
	gnuplot_in_dst = open(tmp_fname1_dst, "w")
except:
	print "Cannot create tmp file %s" % tmp_fname1_dst
	sys.exit(2)

try:
	gnuplot_dst = open(tmp_fname2_dst, "w")
except:
	print "Cannot create tmp file %s" % tmp_fname2_dst
	sys.exit(2)

for port in ports:
	gnuplot_in.write("%d %d\n" % (port, len(hash_port[port].keys())))
gnuplot_in.close()

for dst in dsts:
	gnuplot_in_dst.write("%d %d\n" % (dst, len(hash_dst[dst].keys())))
gnuplot_in_dst.close()

dst_fname = sys.argv[2].replace(".eps", "-dst.eps")

gnuplot.write("""
set terminal postscript eps enhanced color
set output "%s"

set xlabel "Port number"
set ylabel "Number of unique flows"

set logscale x
set logscale y

set xtic auto
set ytic auto

set title "%s"
show title

plot [:][:] '%s' using 1:2 with points notitle
""" % (sys.argv[2], title, tmp_fname1))

fin.close()
gnuplot.close()

os.system("gnuplot " + tmp_fname2)
os.remove(tmp_fname1)
os.remove(tmp_fname2)

gnuplot_dst.write("""
set terminal postscript eps enhanced color
set output "%s"

set xlabel "Destination IP (host address)"
set ylabel "Number of unique source"

set xtic auto
set ytic auto

set title "%s"
show title

plot [:][:] '%s' using 1:2 with points notitle
""" % (dst_fname, title, tmp_fname1_dst))

gnuplot_dst.close()

os.system("gnuplot " + tmp_fname2_dst)
os.remove(tmp_fname1_dst)
os.remove(tmp_fname2_dst)

